from TLSMapper.TLSSUT import *
from TLSMapper.TLSProtocol import *
from datetime import datetime
from aalpy.utils import visualize_automaton
from aalpy.learning_algs import run_Lstar
from aalpy.oracles import RandomWMethodEqOracle, StatePrefixEqOracle
import os
import argparse
import subprocess
import signal
import time
import socket
import logging

from learning.StatePrefixEqOracleFailSafe import StatePrefixOracleFailSafe
from learning.FailSafeCacheSUL import FailSafeCacheSUL
from logger import init_global_logger


# 'OpenSSL', 'WolfSSL', 'MatrixSSL', 'GnuTLS', 'mbedTLS', 'BoringSSL', 'tlse'
# python3 ./model_learning_tlse.py --mode tls12_oneway --auto-start
ciphersuites=[
              CipherSuite.TLS_AES_128_GCM_SHA256,
              CipherSuite.TLS_AES_256_GCM_SHA384,
              CipherSuite.TLS_ECDHE_ECDSA_WITH_AES_256_GCM_SHA384,
              CipherSuite.TLS_DHE_RSA_WITH_CHACHA20_POLY1305_SHA256,
              CipherSuite.TLS_ECDHE_RSA_WITH_AES_256_GCM_SHA384
              ]

# 证书路径配置（写死）
CERT_BASE = '/home/ubuntu/experiments/certs'
CLIENT_KEY = f'{CERT_BASE}/client-key.pem'
CLIENT_CERT = f'{CERT_BASE}/client-cert.pem'
SERVER_CERT = f'{CERT_BASE}/server-cert.pem'
SERVER_KEY = f'{CERT_BASE}/server-key.pem'
CA_CERT = f'{CERT_BASE}/ca-cert.pem'
SERVER_HOST = '127.0.0.1'
SERVER_PORT = 1111

# 服务器路径配置
TLSE_ROOT = '/home/ubuntu/experiments/tlse'
SERVER_BIN = f'{TLSE_ROOT}/build/examples/server'

# TLS 1.2 字母表（不包含 TLS 1.3 专用消息）
TLS12_ALPHABET = [
    'ClientHello',
    'Certificate',
    'ClientKeyExchange',
    'CertificateVerify',
    'ChangeCipherSpec',
    'Finish',
    'ApplicationData',
    'ResumptionClientHello',
    'EmptyCertificate',
    'ClosureAlert'
]

# TLS 1.3 字母表（包含所有 TLS 1.3 支持的消息）
TLS13_ALPHABET = [
    'ClientHello',
    'ClientHelloEmptyKeyShare',  # TLS 1.3 专用：用于 HelloRetryRequest
    'Certificate',
    'CertificateVerify',
    'ChangeCipherSpec',
    'Finish',
    'ApplicationData',
    'ResumptionClientHello',
    'EmptyCertificate',
    'KeyUpdate',  # TLS 1.3 专用：密钥更新
    'ClosureAlert'
]




def get_happy_flow(mode):
    """根据模式返回 happy flow 序列"""
    happy_flows = {
        'tls12_oneway': [
            'ClientHello',
            'ClientKeyExchange',
            'ChangeCipherSpec',
            'Finish',
            'ApplicationData'
        ],
        'tls12_mutual': [
            'ClientHello',
            'Certificate',
            'ClientKeyExchange',
            'CertificateVerify',
            'ChangeCipherSpec',
            'Finish',
            'ApplicationData'
        ],
        'tls13_oneway': [
            'ClientHello',
            'Finish',
            'ApplicationData'
        ],
        'tls13_mutual': [
            'ClientHello',
            'Certificate',
            'CertificateVerify',
            'Finish',
            'ApplicationData'
        ]
    }
    return happy_flows.get(mode, [])


def get_mode_config(mode):
    """根据模式返回配置信息"""
    # 基础服务器命令（单向认证）
    base_server_cmd_oneway = [
        SERVER_BIN,
        '0.0.0.0',
        str(SERVER_PORT),
        SERVER_CERT,
        SERVER_KEY,
        '-f'
    ]
    
    # 双向认证服务器命令（添加 -d 参数）
    base_server_cmd_mutual = [
        SERVER_BIN,
        '0.0.0.0',
        str(SERVER_PORT),
        SERVER_CERT,
        SERVER_KEY,
        '-f',
        '-d'  # -d 代表双向认证
    ]
    
    configs = {
        'tls12_oneway': {
            'version': (3, 3),  # TLS 1.2
            'need_client_cert': False,
            'alphabet': TLS12_ALPHABET,
            'description': 'TLS 1.2 - 单向认证（服务器证书）',
            'server_cmd': base_server_cmd_oneway
        },
        'tls12_mutual': {
            'version': (3, 3),  # TLS 1.2
            'need_client_cert': True,
            'alphabet': TLS12_ALPHABET,
            'description': 'TLS 1.2 - 双向认证（Mutual Authentication）',
            'server_cmd': base_server_cmd_mutual
        },
        'tls13_oneway': {
            'version': (3, 4),  # TLS 1.3
            'need_client_cert': False,
            'alphabet': TLS13_ALPHABET,
            'description': 'TLS 1.3 - 单向认证（服务器证书）',
            'server_cmd': base_server_cmd_oneway
        },
        'tls13_mutual': {
            'version': (3, 4),  # TLS 1.3
            'need_client_cert': True,
            'alphabet': TLS13_ALPHABET,
            'description': 'TLS 1.3 - 双向认证（Mutual Authentication）',
            'server_cmd': base_server_cmd_mutual
        }
    }
    return configs.get(mode)


def create_tls_protocol(mode_config):
    """创建 TLS 协议配置"""
    return TLSProtocol(
        version=mode_config['version'],
        cipher_suite=ciphersuites,
        implementation='tlse',
        verify=False,
        fuzz_tls=(False, None, None),
        target=(SERVER_HOST, SERVER_PORT)
    )


def create_sut(tlsimp, mode_config):
    """创建 SUT（System Under Test）"""
    if mode_config['need_client_cert']:
        return TLSSUT(
            TLSpro=tlsimp,
            keyfile=CLIENT_KEY,
            certfile=CLIENT_CERT,
            server_cert=CA_CERT
        )
    else:
        return TLSSUT(
            TLSpro=tlsimp,
            keyfile=None,
            certfile=None,
            server_cert=CA_CERT
        )


def wait_for_port(host, port, timeout=10):
    """等待端口可用"""
    start_time = time.time()
    attempt = 0
    while time.time() - start_time < timeout:
        try:
            sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            sock.settimeout(1)
            result = sock.connect_ex((host, port))
            sock.close()
            if result == 0:
                elapsed = time.time() - start_time
                if attempt > 0:  # 只在重试后打印
                    print(f"  端口 {host}:{port} 已可用 (等待了 {elapsed:.2f} 秒)")
                return True
        except Exception as e:
            # 只在第一次和每5次尝试后打印错误（避免输出过多）
            if attempt == 0 or attempt % 5 == 0:
                pass  # 静默处理连接错误
        attempt += 1
        time.sleep(0.2)
    
    elapsed = time.time() - start_time
    print(f"  端口 {host}:{port} 在 {elapsed:.2f} 秒内未变为可用")
    return False


def start_server(server_cmd):
    """启动服务器进程"""
    print(f"启动服务器: {' '.join(server_cmd)}")
    proc = subprocess.Popen(
        server_cmd,
        stdout=subprocess.DEVNULL,
        stderr=subprocess.DEVNULL,
        preexec_fn=os.setsid,  # 创建新的进程组
        close_fds=True
    )
    
    # 先等待一小段时间，让服务器进程启动并完成初始化（加载证书等）
    time.sleep(1.0)
    
    # 检查进程是否还在运行（如果立即退出说明启动失败）
    if proc.poll() is not None:
        print(f"✗ 服务器进程已退出 (退出码: {proc.returncode})")
        return None
    
    # 等待服务器启动并监听端口（增加超时时间到15秒）
    print(f"  等待服务器在 {SERVER_HOST}:{SERVER_PORT} 上监听...")
    if wait_for_port(SERVER_HOST, SERVER_PORT, timeout=15):
        print(f"✓ 服务器已启动 (PID: {proc.pid})")
        return proc
    else:
        print("✗ 服务器启动超时")
        # 打印一些调试信息
        print(f"  检查进程状态: {'运行中' if proc.poll() is None else f'已退出 (退出码: {proc.returncode})'}")
        stop_server(proc)
        return None


def stop_server(proc):
    """停止服务器进程"""
    if proc is None:
        return
    try:
        # 终止整个进程组
        os.killpg(os.getpgid(proc.pid), signal.SIGTERM)
        proc.wait(timeout=5)
        print(f"✓ 服务器已停止 (PID: {proc.pid})")
    except ProcessLookupError:
        print(f"服务器进程已不存在 (PID: {proc.pid})")
    except subprocess.TimeoutExpired:
        try:
            os.killpg(os.getpgid(proc.pid), signal.SIGKILL)
            proc.wait()
            print(f"✓ 服务器已强制停止 (PID: {proc.pid})")
        except Exception as e:
            print(f"停止服务器时出错: {e}")
    except Exception as e:
        print(f"停止服务器时出错: {e}")


def test_happy_flow(mode, auto_start_server=False):
    """测试指定模式的 happy flow（正常连接流程）"""
    # 获取模式配置
    mode_config = get_mode_config(mode)
    if not mode_config:
        print(f"错误: 未知的模式 {mode}")
        return False
    
    server_proc = None
    
    try:
        # 自动启动服务器
        if auto_start_server:
            server_proc = start_server(mode_config['server_cmd'])
            if server_proc is None:
                return False
            time.sleep(1)  # 额外等待确保服务器完全就绪
        
        print("\n" + "=" * 80)
        print(f"测试模式: {mode_config['description']}")
        print(f"TLS 版本: {mode_config['version']}")
        print(f"服务器地址: {SERVER_HOST}:{SERVER_PORT}")
        print("=" * 80)
        
        # 创建 TLS 协议配置
        tlsimp = create_tls_protocol(mode_config)
        
        # 创建 SUT
        sul = create_sut(tlsimp, mode_config)
        
        # 获取 happy flow 序列
        happy_flow = get_happy_flow(mode)
        if not happy_flow:
            print(f"错误: 未知的模式 {mode}")
            return False
        
        print(f"\n执行 Happy Flow 序列: {' -> '.join(happy_flow)}")
        print("-" * 80)
        
        # 执行测试
        out, sent_data, received_data = sul.query_to_test(happy_flow)
        
        print("\n✓ Happy Flow 测试成功！")
        print(f"服务器响应序列: {out}")
        print("=" * 80)
        
        return True
        
    except Exception as e:
        print(f"\n✗ Happy Flow 测试失败: {e}")
        import traceback
        traceback.print_exc()
        return False
        
    finally:
        # 停止服务器
        if auto_start_server and server_proc:
            stop_server(server_proc)
            time.sleep(1)  # 等待端口释放


def learn_mode(mode, auto_start_server=False):
    """学习指定模式的状态机"""
    # 获取模式配置
    mode_config = get_mode_config(mode)
    if not mode_config:
        print(f"错误: 未知的模式 {mode}")
        return False
    
    # 初始化日志（保存到文件）
    log_dir = "logs"
    os.makedirs(log_dir, exist_ok=True)
    log_file = f"{log_dir}/model_learning_tlse_{mode}_{datetime.now().strftime('%Y%m%d_%H%M%S')}.log"
    init_global_logger(log_file=log_file, level=logging.INFO, enable_console=True)
    print(f"日志文件: {log_file}")
    
    server_proc = None
    
    try:
        # 自动启动服务器
        if auto_start_server:
            server_proc = start_server(mode_config['server_cmd'])
            if server_proc is None:
                return False
            time.sleep(1)  # 额外等待确保服务器完全就绪
        
        print("\n" + "=" * 80)
        print(f"模式: {mode_config['description']}")
        print(f"TLS 版本: {mode_config['version']}")
        print(f"需要客户端证书: {mode_config['need_client_cert']}")
        print(f"字母表: {mode_config['alphabet']}")
        print(f"服务器地址: {SERVER_HOST}:{SERVER_PORT}")
        print("=" * 80)
        
        # 创建 TLS 协议配置
        tlsimp = create_tls_protocol(mode_config)
        
        # 创建 SUT
        sul = create_sut(tlsimp, mode_config)
        
        # 确保数据库目录存在
        db_dir = "/tmp/db"
        os.makedirs(db_dir, exist_ok=True)
        db_path = f"{db_dir}/tlse_{mode}.db"
        # 创建包装类（用于 L* 算法）
        sul = FailSafeCacheSUL(sul, database=db_path)
        
        # 获取字母表
        alphabet = mode_config['alphabet']
        
        print(f"\n开始学习状态机...")
        print(f"字母表大小: {len(alphabet)}")
        print("-" * 80)
        
        # 获取并执行 happy flow 作为初始查询
        happy_flow = get_happy_flow(mode)
        initial_samples = None
        if happy_flow:
            print(f"\n执行 Happy Flow 作为初始查询: {' -> '.join(happy_flow)}")
            try:
                # 执行 happy flow 查询，让 L* 算法在初始观察表中包含这些信息
                happy_flow_output = sul.query(happy_flow)
                print(f"Happy Flow 响应: {happy_flow_output}")
                # 将 happy flow 作为初始样本传递给 L* 算法
                # samples 格式: [(input_sequence, output_sequence), ...]
                initial_samples = [(happy_flow, happy_flow_output)]
                print(f"已将 Happy Flow 添加到初始样本中")
                print("-" * 80)
            except Exception as e:
                print(f"警告: Happy Flow 查询失败: {e}")
                print("继续学习过程（不使用初始样本）...")
        
        # 创建等价性预言
        eq_oracle = StatePrefixOracleFailSafe(alphabet, sul, walks_per_state=1, walk_len=3, database=db_path)
        eq_oracle.add_manual_counterexample(happy_flow)

        # eq_oracle = StatePrefixEqOracle(alphabet, sul, walks_per_state=300, walk_len=10)
        
        # 生成输出目录
        imp = tlsimp.implementation
        ver = tlsimp.imp_ver if hasattr(tlsimp, 'imp_ver') else 'None'
        if tlsimp.version == (3, 4):
            tv = 'tls1.3'
        else:
            tv = 'tls1.2'
        
        auth_type = 'mutual' if mode_config['need_client_cert'] else 'oneway'
        folder_name = f"{imp}_{ver}_{tv}_{auth_type}"
        output_base_dir = "dots/tlse"
        os.makedirs(output_base_dir, exist_ok=True)
        folder_path = f"{output_base_dir}/{folder_name}"
        os.makedirs(folder_path, exist_ok=True)
        save_dir = f"{folder_path}/model.dot"
        
        print(f"输出目录: {folder_path}")
        print(f"模型文件: {save_dir}")
        print("-" * 80)
        
        # 运行 L* 算法学习状态机
        print("\n运行 L* 算法...")
        # 如果提供了初始样本（happy flow），则传递给 L* 算法
        model = run_Lstar(
            alphabet, 
            sul, 
            eq_oracle, 
            'mealy', 
            samples=initial_samples,  # 传递初始样本（happy flow）
            cache_and_non_det_check=False
        )
        
        # 保存模型
        model.save(file_path=save_dir)
        print(f"\n✓ 模型已保存到: {save_dir}")
        
        # 可视化状态机
        print("\n生成状态机可视化...")
        visualize_automaton(model)

        print("\n" + "=" * 80)
        print(f"✓ {mode_config['description']} 状态机学习完成！")
        print("=" * 80)
        
        return True
        
    except Exception as e:
        print(f"\n✗ 学习过程中出错: {e}")
        import traceback
        traceback.print_exc()
        return False
        
    finally:
        # 停止服务器
        if auto_start_server and server_proc:
            stop_server(server_proc)
            time.sleep(1)  # 等待端口释放


def main():
    parser = argparse.ArgumentParser(
        description='TLSE 服务器四种模式的状态机学习',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
测试模式说明:
  tls12_oneway  - TLS 1.2 单向认证（服务器证书）
  tls12_mutual  - TLS 1.2 双向认证（Mutual Authentication）
  tls13_oneway  - TLS 1.3 单向认证（服务器证书）
  tls13_mutual  - TLS 1.3 双向认证（Mutual Authentication）

使用示例:
  # 学习单个模式（需要手动启动服务器）
  python3 model_learning_tlse.py --mode tls12_oneway
  
  # 自动启动服务器并学习单个模式
  python3 model_learning_tlse.py --mode tls12_oneway --auto-start
  
  # 自动学习所有四种模式（自动启动和停止服务器）
  python3 model_learning_tlse.py --all
  
  # 测试 happy flow（正常连接流程）
  python3 model_learning_tlse.py --mode tls12_oneway --test --auto-start
        """
    )
    
    parser.add_argument(
        '--mode',
        type=str,
        choices=['tls12_oneway', 'tls12_mutual', 'tls13_oneway', 'tls13_mutual'],
        help='选择学习模式（与 --all 互斥）'
    )
    
    parser.add_argument(
        '--all',
        action='store_true',
        help='自动学习所有四种模式'
    )
    
    parser.add_argument(
        '--auto-start',
        action='store_true',
        help='自动启动和停止服务器（仅在使用 --mode 时有效）'
    )
    
    parser.add_argument(
        '--test',
        action='store_true',
        help='测试 happy flow（正常连接流程），需要与 --mode 一起使用'
    )
    
    args = parser.parse_args()
    
    # 验证参数
    if args.all and args.mode:
        print("错误: --all 和 --mode 不能同时使用")
        return
    
    if not args.all and not args.mode:
        print("错误: 必须指定 --mode 或 --all")
        return
    
    if args.test and not args.mode:
        print("错误: --test 必须与 --mode 一起使用")
        return
    
    # 学习所有模式
    if args.all:
        print("\n" + "=" * 80)
        print("自动学习所有四种模式的状态机")
        print("=" * 80)
        
        modes = ['tls12_oneway', 'tls12_mutual', 'tls13_oneway', 'tls13_mutual']
        results = {}
        
        for i, mode in enumerate(modes, 1):
            print(f"\n\n{'#' * 80}")
            print(f"# 模式 {i}/4: {mode}")
            print(f"{'#' * 80}\n")
            
            success = learn_mode(mode, auto_start_server=True)
            results[mode] = success
            
            if i < len(modes):
                print("\n等待 2 秒后继续下一个模式...")
                time.sleep(2)
        
        # 打印总结
        print("\n\n" + "=" * 80)
        print("学习总结")
        print("=" * 80)
        for mode, success in results.items():
            status = "✓ 成功" if success else "✗ 失败"
            mode_config = get_mode_config(mode)
            if mode_config:
                print(f"{mode_config['description']}: {status}")
        print("=" * 80)
        
    else:
        # 如果是测试模式
        if args.test:
            test_happy_flow(args.mode, auto_start_server=args.auto_start)
        else:
            # 学习单个模式
            learn_mode(args.mode, auto_start_server=args.auto_start)


if __name__ == '__main__':
    main()

